// instrument.go - 一个单文件的 Go 函数调用追踪插桩工具。
// 它通过修改 Go 源代码的抽象语法树 (AST) 来实现，
// 在函数入口处自动插入 `defer Trace()()` 语句。
package main

import (
	"bytes"         // 用于高效处理字节序列，特别是格式化代码输出
	"flag"          // 用于处理命令行参数（如 -w）
	"fmt"           // 用于格式化输入输出，包括打印追踪信息和错误信息
	"go/ast"        // Go 语言抽象语法树相关类型和函数
	"go/format"     // 用于格式化 Go 源代码
	"go/parser"     // 用于将 Go 源代码解析为 AST
	"go/token"      // Go 语言词法单元和文件集（管理源码位置）
	"io/fs"         // Go 1.16+ 文件系统接口，用于 WalkDir 和文件权限
	"os"            // 操作系统功能，如文件操作、命令行参数、环境变量
	"path/filepath" // 用于处理和操作文件路径，跨平台兼容
	"runtime"       // 提供与 Go 运行时交互的功能，如获取调用栈、goroutine ID
	"strconv"       // 字符串和其他基本数据类型之间的转换
	"strings"       // 字符串操作函数
	"sync"          // 提供基本的同步原语，如 sync.Map

	// **关键修复点:** 重新引入并使用 astutil.Apply 进行 AST 修改
	"golang.org/x/tools/go/ast/astutil" // AST 操作的辅助函数 (推荐用于修改)
)

// ============================================================================
// 1. 追踪逻辑 (实现具体的追踪行为)
// ============================================================================

var (
	// goroutineSpace 用于在 runtime.Stack 输出中识别 goroutine ID 的前缀字节序列。
	goroutineSpace = []byte("goroutine ")
	// goroutineIndent 使用 sync.Map 存储每个 goroutine 的当前调用栈深度（用于缩进）。
	// sync.Map 适用于读多写少且键值对数量不定的场景，适合存储 goroutine 状态。
	// key: goroutine ID (uint64), value: indent level (int64)
	goroutineIndent sync.Map
)

// 定义 ANSI 转义序列颜色常量, 用于在支持的终端中为输出着色，提高可读性。
const (
	colorReset   = "\033[0m"  // 重置所有颜色和样式
	colorGreen   = "\033[32m" // 绿色，用于 Goroutine ID
	colorBlue    = "\033[34m" // 蓝色，用于函数名
	colorMagenta = "\033[35m" // 品红色，用于行号
	// traceFuncName 是我们注入和识别的追踪函数的名字。
	// 在单文件工具中硬编码此名称简化了 AST 检查逻辑。
	traceFuncName = "Trace"
)

// curGoroutineID 获取当前执行的 goroutine 的 ID。
// 注意：此方法依赖 runtime.Stack 输出的特定格式，这种格式在不同 Go 版本间可能变化，
// 并且获取堆栈信息有性能开销。它主要用于调试和演示目的。
func curGoroutineID() uint64 {
	var buf [64]byte                                     // 在栈上分配一个小缓冲区，避免堆分配
	n := runtime.Stack(buf[:], false)                    // 获取当前 goroutine 的堆栈跟踪信息
	idField := bytes.TrimPrefix(buf[:n], goroutineSpace) // 移除 "goroutine " 前缀
	i := bytes.IndexByte(idField, ' ')                   // 查找 ID 后面的第一个空格
	if i < 0 {
		fmt.Fprintf(os.Stderr, "警告: 获取 goroutine ID 时无法解析堆栈: %q\n", buf[:n])
		return 0 // 返回 0 表示获取失败
	}
	id, err := strconv.ParseUint(string(idField[:i]), 10, 64) // 解析 ID
	if err != nil {
		fmt.Fprintf(os.Stderr, "警告: 解析 goroutine ID '%s' 失败: %v\n", string(idField[:i]), err)
		return 0 // 返回 0 表示解析失败
	}
	return id
}

// printTrace 格式化并打印单条追踪日志（函数进入或退出）。
func printTrace(id uint64, name string, line int, arrow string, indent int) {
	indents := strings.Repeat("  ", indent) // 使用两个空格作为一级缩进
	fmt.Printf("%sgid[%05d]%s %s%s %s%s%s:%s%d%s\n",
		colorGreen, id, colorReset, // Goroutine ID
		indents, arrow, // 缩进和箭头
		colorBlue, name, colorReset, // 函数名
		colorMagenta, line, colorReset) // 行号
}

// Trace 是实际被注入到目标 Go 代码中的追踪函数。
func Trace() func() {
	pc, _, line, ok := runtime.Caller(1) // 获取调用者信息
	if !ok {
		fmt.Fprintf(os.Stderr, "警告: runtime.Caller(1) 失败\n")
		return func() {}
	}
	fn := runtime.FuncForPC(pc) // 获取函数信息
	if fn == nil {
		fmt.Fprintf(os.Stderr, "警告: runtime.FuncForPC(%v) 失败\n", pc)
		return func() {}
	}
	name := fn.Name()       // 获取函数名
	gid := curGoroutineID() // 获取 Goroutine ID

	// --- 函数进入 ---
	rawIndent, _ := goroutineIndent.LoadOrStore(gid, int64(0)) // 获取或设置缩进
	indent := rawIndent.(int64)
	printTrace(gid, name, line, "->", int(indent)) // 打印进入日志
	goroutineIndent.Store(gid, indent+1)           // 增加缩进

	// --- 返回退出时执行的闭包 ---
	return func() {
		// --- 函数退出 ---
		currentRawIndent, loaded := goroutineIndent.Load(gid) // 加载当前缩进
		if !loaded {
			fmt.Fprintf(os.Stderr, "警告: goroutine %d 的缩进信息在退出时意外丢失 (函数: %s)\n", gid, name)
			printTrace(gid, name, line, "<-", int(indent)) // 使用进入时的缩进回退
			return
		}
		currentIndent := currentRawIndent.(int64)
		printTrace(gid, name, line, "<-", int(indent)) // 使用进入时的缩进打印退出日志
		newExitIndent := currentIndent - 1             // 计算新缩进
		if newExitIndent <= 0 {
			goroutineIndent.Delete(gid) // 缩进为0则移除
		} else {
			goroutineIndent.Store(gid, newExitIndent) // 更新缩进
		}
	}
}

// ============================================================================
// 2. AST 插桩逻辑 (使用 astutil.Apply)
// ============================================================================

// instrumenter 结构体封装了 AST 修改的相关逻辑。
type instrumenter struct{}

// NewInstrumenter 创建一个新的 instrumenter 实例。
func NewInstrumenter() *instrumenter {
	return &instrumenter{}
}

// Instrument 是核心的插桩处理函数。
// **关键修复点:** 使用 astutil.Apply 进行 AST 遍历和修改，而不是 ast.Inspect。
func (a *instrumenter) Instrument(filename string) ([]byte, error) {
	fset := token.NewFileSet()
	curAST, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
	if err != nil {
		return nil, fmt.Errorf("解析文件 %s 时出错: %w", filename, err)
	}

	var modified bool // 标记是否进行了修改

	// 使用 astutil.Apply 来遍历 AST 节点。
	// Apply 提供了 Cursor，更适合在遍历时安全地修改 AST。
	// pre 函数在访问节点前调用，post 函数在访问节点后调用。
	astutil.Apply(curAST, func(c *astutil.Cursor) bool {
		// 检查当前节点是否为函数声明
		fd, ok := c.Node().(*ast.FuncDecl)
		// 如果不是函数声明或没有函数体，则继续遍历子节点
		if !ok || fd.Body == nil {
			return true // true 表示继续遍历
		}
		// 尝试向函数体添加 defer 语句
		if a.addDeferStmt(fd) {
			modified = true // 标记已修改
			// 注意：虽然我们修改了 fd.Body.List，但 Apply 会处理好后续遍历
		}
		return true // 继续遍历子节点和兄弟节点
	}, nil) // post 函数为 nil，我们不需要在访问后操作

	// 如果未进行任何修改，返回 nil
	if !modified {
		return nil, nil
	}

	// 格式化修改后的 AST
	var buf bytes.Buffer
	err = format.Node(&buf, fset, curAST)
	if err != nil {
		return nil, fmt.Errorf("格式化 %s 的新代码时出错: %w", filename, err)
	}
	return buf.Bytes(), nil
}

// addDeferStmt 尝试向给定的函数声明 (fd) 的函数体开头插入 "defer Trace()()" 语句。
// (此函数逻辑不变，之前的实现是正确的)
func (a *instrumenter) addDeferStmt(fd *ast.FuncDecl) (added bool) {
	if fd.Body == nil {
		return false
	}
	if fd.Body.List == nil {
		fd.Body.List = make([]ast.Stmt, 0, 1)
	}
	stmts := fd.Body.List

	// 幂等性检查
	for _, stmt := range stmts {
		ds, ok := stmt.(*ast.DeferStmt)
		if !ok {
			continue
		}
		outerCall, ok := ds.Call.Fun.(*ast.CallExpr)
		if !ok {
			continue
		}
		ident, ok := outerCall.Fun.(*ast.Ident)
		if !ok {
			continue
		}
		if ident.Name == traceFuncName {
			return false // 已存在
		}
	}

	// 构造新节点
	newDeferStmt := &ast.DeferStmt{
		Call: &ast.CallExpr{
			Fun: &ast.CallExpr{
				Fun: &ast.Ident{Name: traceFuncName},
			},
		},
	}

	// 插入到开头
	newList := make([]ast.Stmt, 0, len(stmts)+1)
	newList = append(newList, newDeferStmt)
	newList = append(newList, stmts...)
	fd.Body.List = newList

	return true // 成功添加
}

// ============================================================================
// 3. 命令行处理逻辑 (无变化)
// ============================================================================

var (
	// write 标志决定是否将修改写回源文件
	write = flag.Bool("w", false, "将结果写入源文件而不是标准输出")
)

// usage 打印命令行用法信息
func usage() {
	fmt.Fprintf(os.Stderr, "用法: go run instrument.go [-w] [path]\n")
	fmt.Fprintf(os.Stderr, "       instrument [-w] [path] (编译后)\n\n")
	fmt.Fprintf(os.Stderr, "  path: 要处理的 Go 文件或目录。\n")
	fmt.Fprintf(os.Stderr, "        默认当前目录。递归处理 .go 文件 (跳过 _test.go, vendor/, .* 目录)。\n")
	fmt.Fprintf(os.Stderr, "选项:\n")
	flag.PrintDefaults()
}

// main 函数是程序的入口点
func main() {
	flag.Usage = usage
	flag.Parse()

	args := flag.Args()
	targetPath := "."
	if len(args) > 0 {
		targetPath = args[0]
	}

	absTargetPath, err := filepath.Abs(targetPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "错误: 获取绝对路径 '%s' 失败: %v\n", targetPath, err)
		os.Exit(1)
	}

	targetInfo, err := os.Stat(absTargetPath)
	if err != nil {
		if os.IsNotExist(err) {
			fmt.Fprintf(os.Stderr, "错误: 路径 '%s' 不存在\n", absTargetPath)
		} else {
			fmt.Fprintf(os.Stderr, "错误: 访问路径 '%s' 时出错: %v\n", absTargetPath, err)
		}
		os.Exit(1)
	}

	ins := NewInstrumenter()

	// --- 获取工具自身路径 ---
	selfPath := ""
	if execPath, execErr := os.Executable(); execErr == nil {
		selfPath, _ = filepath.Abs(execPath)
	}
	runFilePath := ""
	_, runFile, _, runOk := runtime.Caller(0)
	if runOk {
		runFilePath, _ = filepath.Abs(runFile)
	}
	// --- 工具路径获取结束 ---

	stats := struct{ processed, modified, skipped, errors int }{}

	if targetInfo.IsDir() {
		fmt.Printf("递归处理目录: %s\n", absTargetPath)
		err := filepath.WalkDir(absTargetPath, func(path string, d fs.DirEntry, walkErr error) error {
			// --- 文件/目录过滤逻辑 ---
			if walkErr != nil {
				stats.errors++
				return nil
			}
			absPath, pathErr := filepath.Abs(path)
			if pathErr != nil {
				stats.errors++
				return nil
			}
			if d.IsDir() {
				dirName := d.Name()
				if dirName == "vendor" || (dirName != "." && dirName != ".." && strings.HasPrefix(dirName, ".")) {
					stats.skipped++
					return filepath.SkipDir
				}
				return nil
			}
			if !d.Type().IsRegular() {
				stats.skipped++
				return nil
			}
			fileName := d.Name()
			if !strings.HasSuffix(fileName, ".go") || strings.HasSuffix(fileName, "_test.go") {
				if strings.HasSuffix(fileName, ".go") {
					stats.skipped++
				}
				return nil
			}
			if absPath == selfPath || absPath == runFilePath {
				stats.skipped++
				return nil
			}
			// --- 过滤结束 ---

			// --- 处理文件 ---
			fmt.Printf("处理文件: %s\n", path)
			stats.processed++
			modified, procErr := processFile(path, ins, *write)
			if procErr != nil {
				stats.errors++
			} else if modified {
				stats.modified++
			}
			return nil
		})
		if err != nil {
			stats.errors++
		}
	} else {
		// --- 处理单个文件 ---
		fileName := targetInfo.Name()
		absPath := absTargetPath
		// --- 文件过滤 ---
		if !strings.HasSuffix(fileName, ".go") || strings.HasSuffix(fileName, "_test.go") {
			if strings.HasSuffix(fileName, "_test.go") {
				stats.skipped++
			} else {
				stats.errors++
			}
		} else if absPath == selfPath || absPath == runFilePath {
			stats.skipped++
		} else {
			// --- 过滤结束 ---
			// --- 处理文件 ---
			fmt.Printf("处理单个文件: %s\n", absPath)
			stats.processed++
			modified, procErr := processFile(absPath, ins, *write)
			if procErr != nil {
				stats.errors++
			} else if modified {
				stats.modified++
			}
		}
	}

	// --- 打印总结 ---
	fmt.Println("\n--- 处理完成 ---")
	fmt.Printf("已处理 Go 文件数: %d\n", stats.processed)
	fmt.Printf("成功修改文件数:  %d\n", stats.modified)
	fmt.Printf("跳过文件/目录数: %d\n", stats.skipped)
	fmt.Printf("处理失败文件数:  %d\n", stats.errors)

	if stats.errors > 0 {
		os.Exit(1)
	}
}

// processFile 函数负责处理单个文件的插桩操作，包括可选的文件写入。
// (此函数逻辑不变)
func processFile(path string, ins *instrumenter, writeToFile bool) (bool, error) {
	newSrc, err := ins.Instrument(path)
	if err != nil {
		return false, fmt.Errorf("插桩分析失败: %w", err)
	}
	if newSrc == nil {
		return false, nil // 未修改
	}

	modified := true // 已修改
	if !writeToFile {
		fmt.Println("  (检测到修改，但未写入文件。使用 -w 选项写入)")
	} else {
		originalInfo, statErr := os.Stat(path)
		perms := fs.FileMode(0644) // 默认权限
		if statErr == nil {
			perms = originalInfo.Mode().Perm()
		} else {
			fmt.Fprintf(os.Stderr, "警告: 无法获取文件 %s 的原始权限: %v, 将使用默认权限 0644\n", path, statErr)
		}
		err = os.WriteFile(path, newSrc, perms)
		if err != nil {
			return modified, fmt.Errorf("写入文件 '%s' 失败: %w", path, err)
		}
	}
	return modified, nil
}
